import numpy as np
import mindspore, time, os,sys
sys.path.append(os.path.abspath(os.path.join(os.getcwd())))
path = os.path.abspath(os.path.join(os.getcwd()))
from mindspore import context, nn, ops, Tensor, save_checkpoint
from mindspore.context import ParallelMode
from mindspore.train.callback import Callback
from mind3d.utils.PointTransformerUtils import to_categorical, pointnet_to_categorical


def get_rank_id():
    global_rank_id = os.getenv('RANK_ID', '0')
    return int(global_rank_id)


class ModifyNetwork(nn.Cell):
    def __init__(self, models):
        super(ModifyNetwork, self).__init__()
        self.models = models
    def construct(self, data):
        output = self.models(data)
        return output

class CustumWithLoss(nn.Cell):
    def __init__(self, models, loss_fn, num_part):
        super(CustumWithLoss, self).__init__()
        self.models = models
        self.Loss = loss_fn
        self.num_part = num_part
        self.cat = ops.Concat(axis = -1)

    def construct(self, point_set, cls, target):
        _,N,_ = point_set.shape
        #print(mindspore.numpy.tile(to_categorical(cls, 16), (1, N, 1)).shape)
        point = self.cat((point_set, mindspore.numpy.tile(pointnet_to_categorical(cls, 16), (1, N, 1))))
        output = self.models(point)
        loss = self.Loss(output.view(-1, self.num_part), target.view(-1, 1)[:, 0])
        return loss
    

class WithEvalCell(nn.Cell):
    def __init__(self, network, add_cast_fp32=False):
        super(WithEvalCell, self).__init__(auto_prefix=False)
        assert isinstance(add_cast_fp32, bool), "add_cast_fp32 should be float"
        self.Network = network
        self.cat = ops.Concat(axis=-1)

    def construct(self, point_set, cls, target):
        _, N, _ = point_set.shape
        point = self.cat((point_set, mindspore.numpy.tile(pointnet_to_categorical(cls, 16), (1, N, 1))))
        pred = self.Network(point)
        return  pred, target


class pointnetCustomWithLossCell(nn.Cell):
    def __init__(self, model, loss_fn, mode):
        super(pointnetCustomWithLossCell, self).__init__(auto_prefix=False)
        self.model = model
        self.loss = loss_fn
        self.mode = mode

    def construct(self, point_set, cls, seg_label):  # B N C, B 1, B N
        output = self.model(point_set, pointnet_to_categorical(cls, 16)) 
        loss = self.loss(output, seg_label, method=self.mode)
        return loss
    

class pointnetWithEvalCell(nn.Cell):
    def __init__(self, network):
        super(pointnetWithEvalCell, self).__init__(auto_prefix=False)
        self._network = network

    def construct(self, point_set, cls, target):
        pred = self._network(point_set, pointnet_to_categorical(cls, 16))
        return  pred, target


classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
           'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37],
           'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49],
           'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

label_to_cat = {}  # {0:Airplane, 1:Airplane, ...49:Table}
for cls in classes:
    for label in classes[cls]:
        label_to_cat[label] = cls

class IoU(nn.Metric):
    def __init__(self):
        super(IoU, self).__init__()
        self.clear()

    def clear(self):
        """Clears the internal evaluation result."""
        self.shape_ious = {cat: [] for cat in classes}

    def update(self, *inputs):
        if len(inputs) != 2:
            raise ValueError('Distribute accuracy needs 2 input (y_correct), but got {}'.format(len(inputs)))
        y_pred = self._convert_data(inputs[0])
        target = self._convert_data(inputs[1])
        B, N, _ = y_pred.shape
        cur_pred_val_logits = y_pred
        cur_pred_val = np.zeros((B, N)).astype(np.int32)

        for i in range(B):
            cat = label_to_cat[target[i, 0]]
            logits = cur_pred_val_logits[i, :, :]
            cat_logits = logits[:, classes[cat]]
            cur_pred_val[i, :] = np.argmax(cat_logits, 1) + classes[cat][0]

        for i in range(B):
            pred_seg = cur_pred_val[i, :]
            leb_seg = target[i, :]
            cat = label_to_cat[leb_seg[0]]
            part_ious = [0.0 for _ in range(len(classes[cat]))]
            for category in classes[cat]:
                if (np.sum(leb_seg == category) == 0) and (
                        np.sum(pred_seg == category) == 0):
                    part_ious[category - classes[cat][0]] = 1.0
                else:
                    part_ious[category - classes[cat][0]] = np.sum((leb_seg == category) & (pred_seg == category)) / \
                                                       float(np.sum((leb_seg == category) | (pred_seg == category)))
            self.shape_ious[cat].append(np.mean(part_ious))


    def eval(self):
        all_shape_ious = []
        class_shape_ious = {}
        for category in self.shape_ious:
            for iou in self.shape_ious[category]:
                all_shape_ious.append(iou)
            class_shape_ious[category] = np.mean(self.shape_ious[category])

        class_avg_iou = np.mean(list(class_shape_ious.values()))
        inctance_avg_iou = np.mean(all_shape_ious)

        return class_avg_iou, inctance_avg_iou
    

class CheckLoss(Callback):
    """
    Monitor the loss in training.

    If the loss is NAN or INF, it will terminate training.

    Raises:
        ValueError: If per_print_steps is not an integer or less than zero.
    """

    def begin(self, run_context):
        self.train_time = time.time()
        if get_rank_id() == 0:
            print('Train start at {}'.format(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())))

    def step_end(self, run_context):
        """
        Print training loss at the end of step.

        Args:
            run_context (RunContext): Context of the train running.
        """
        callback_params = run_context.original_args()
        loss = callback_params.net_outputs
        num_epoch = callback_params.cur_epoch_num

        if isinstance(loss, (tuple, list)):
            if isinstance(loss[0], Tensor) and isinstance(loss[0].asnumpy(), np.ndarray):
                loss = loss[0]

        if isinstance(loss, Tensor) and isinstance(loss.asnumpy(), np.ndarray):
            loss = float(np.mean(loss.asnumpy()))

        step_in_epoch = (callback_params.cur_step_num - 1) % callback_params.batch_num + 1

        if isinstance(loss, float) and (np.isnan(loss) or np.isinf(loss)):
            raise ValueError("epoch: {} step: {}. Invalid loss, terminating training.".format(
                callback_params.cur_epoch_num, step_in_epoch))

        rank_id = get_rank_id()

        print(f"Device:{rank_id}, Epoce:{num_epoch}, Step:{step_in_epoch}, Train loss:{loss:.3f}", flush=True)

    def end(self, run_context):
        training_time = int(time.time() - self.train_time)
        if get_rank_id() == 0:
            print('Train end at {}'.format(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())))
            print(f"Training time {training_time//3600}h{(training_time%3600)//60}m{(training_time%3600)%60}s")

    
class CallbackSaveByIoU(Callback):
    """SaveCallback"""
    def __init__(self, eval_model, ds_eval, eval_period=1, eval_start=1, save_path=None):
        """init"""
        super(CallbackSaveByIoU, self).__init__()
        self.model = eval_model
        self.ds_eval = ds_eval
        self.ins_mIoU = 0.
        self.cls_mIoU = 0.
        self.eval_period = eval_period
        self.save_path = save_path
        self.eval_start = eval_start

    def epoch_end(self, run_context):
        """epoch end"""
        cb_params = run_context.original_args()
        cur_epoch = cb_params.cur_epoch_num
        rank_id = get_rank_id()
        if ((cur_epoch + 1) % self.eval_period) == 0:
            if cur_epoch < self.eval_start:
                return
            if rank_id == 0:
                print("Start evaluate...")
            result = self.model.eval(self.ds_eval)
            cls_mIoU = result['IoU'][0]
            ins_mIoU = result['IoU'][1]
            if cls_mIoU > self.cls_mIoU:
                self.cls_mIoU = cls_mIoU
            if ins_mIoU > self.ins_mIoU:
                self.ins_mIoU = ins_mIoU
                file_name = f"best_model_dev_{rank_id}.ckpt"
                save_path = os.path.join(self.save_path, file_name)
                print("Save model...")
                save_checkpoint(save_obj=cb_params.train_network, ckpt_file_name=save_path)
            print(f"Device:{rank_id}, Epoce:{cur_epoch}, Instance mIoU:{ins_mIoU:.5f}, Class mIoU:{cls_mIoU:.5f}")

    def end(self, run_context):
        _ = run_context.original_args()
        rank_id = get_rank_id()
        print(f"Device:{rank_id}, Best Instance mIoU:{(self.ins_mIoU*100):.2f}%, Class mIoU:{(self.cls_mIoU*100):.2f}%")

class mIoU(nn.Metric):
    '''compute miou'''
    def __init__(self, test_steps_per_epoch, batch_size, num_classes):
        super(mIoU, self).__init__()
        self.test_steps_per_epoch =test_steps_per_epoch
        self.batch_size = batch_size
        self.num_classes =num_classes

        self.clear()

    def clear(self):
        """Clears the internal evaluation result."""
        self.shape_ious = {cat: [] for cat in classes}
        self.pointacc_list = []
        self.pointacc_per_class_array = np.zeros((self.test_steps_per_epoch, self.num_classes))
        self.voxacc_list = []
        self.voxacc_per_class_array = np.zeros((self.test_steps_per_epoch, self.num_classes))
        self.voxcaliacc_list = []
        self.pointmiou_per_class_array = np.zeros((self.test_steps_per_epoch, self.num_classes))
        self.voxmiou_per_class_array = np.zeros((self.test_steps_per_epoch, self.num_classes))
        self.masks = np.zeros((self.test_steps_per_epoch, self.num_classes))
        self.id = 0

    def filter_points(self, coords, preds, targets, weights):
        """filter points"""
        assert coords.shape[0] == preds.shape[0] == targets.shape[0] == weights.shape[0]
        coord_hash = [hash(str(coords[point_idx][0]) + str(coords[point_idx][1]) + str(coords[point_idx][2])) for
                      point_idx in range(coords.shape[0])]
        _, coord_ids = np.unique(np.array(coord_hash), return_index=True)
        coord_filtered, pred_filtered, target_filtered, weight_filtered = coords[coord_ids], preds[coord_ids], targets[
            coord_ids], weights[coord_ids]

        return coord_filtered, pred_filtered, target_filtered, weight_filtered

    def point_cloud_label_to_surface_voxel_label_fast(self, point_cloud, label, res=0.0484):
        """Convert point cloud label to surface voxel label"""
        coordmax = np.max(point_cloud, axis=0)
        coordmin = np.min(point_cloud, axis=0)
        nvox = np.ceil((coordmax - coordmin) / res)
        vidx = np.ceil((point_cloud - coordmin) / res)
        vidx = vidx[:, 0] + vidx[:, 1] * nvox[0] + vidx[:, 2] * nvox[0] * nvox[1]
        uvidx, vpidx = np.unique(vidx, return_index=True)
        if label.ndim == 1:
            uvlabel = label[vpidx]
        else:
            assert label.ndim == 2
        uvlabel = label[vpidx, :]
        return uvidx, uvlabel, nvox

    def compute_acc(self, coords, preds, targets, weights, num_classes):
        """compute acc"""
        coords, preds, targets, weights = self.filter_points(coords, preds, targets, weights)

        seen_classes = np.unique(targets)
        mask = np.zeros(num_classes)
        mask[seen_classes] = 1

        total_correct = 0
        total_seen = 0
        total_seen_class = [0 for _ in range(num_classes)]
        total_correct_class = [0 for _ in range(num_classes)]

        total_correct_vox = 0
        total_seen_vox = 0
        total_seen_class_vox = [0 for _ in range(num_classes)]
        total_correct_class_vox = [0 for _ in range(num_classes)]

        labelweights = np.zeros(num_classes)
        labelweights_vox = np.zeros(num_classes)

        correct = np.sum(preds == targets)  # evaluate only on 20 categories but not unknown
        total_correct += correct
        total_seen += targets.shape[0]
        tmp, _ = np.histogram(targets, range(num_classes + 1))
        labelweights += tmp
        for l in seen_classes:
            total_seen_class[l] += np.sum(targets == l)
            total_correct_class[l] += np.sum((preds == l) & (targets == l))

        _, uvlabel, _ = self.point_cloud_label_to_surface_voxel_label_fast(coords, np.concatenate(
            (np.expand_dims(targets, 1), np.expand_dims(preds, 1)), axis=1), res=0.02)
        total_correct_vox += np.sum(uvlabel[:, 0] == uvlabel[:, 1])
        total_seen_vox += uvlabel[:, 0].shape[0]
        tmp, _ = np.histogram(uvlabel[:, 0], range(num_classes + 1))
        labelweights_vox += tmp
        for l in seen_classes:
            total_seen_class_vox[l] += np.sum(uvlabel[:, 0] == l)
            total_correct_class_vox[l] += np.sum((uvlabel[:, 0] == l) & (uvlabel[:, 1] == l))

        pointacc = total_correct / float(total_seen)
        voxacc = total_correct_vox / float(total_seen_vox)

        labelweights = labelweights.astype(np.float32) / np.sum(labelweights.astype(np.float32))
        labelweights_vox = labelweights_vox.astype(np.float32) / np.sum(labelweights_vox.astype(np.float32))
        caliweights = labelweights_vox
        voxcaliacc = np.average(
            np.array(total_correct_class_vox) / (np.array(total_seen_class_vox, dtype=np.float64) + 1e-8),
            weights=caliweights)

        pointacc_per_class = np.zeros(num_classes)
        voxacc_per_class = np.zeros(num_classes)
        for l in seen_classes:
            pointacc_per_class[l] = total_correct_class[l] / (total_seen_class[l] + 1e-8)
            voxacc_per_class[l] = total_correct_class_vox[l] / (total_seen_class_vox[l] + 1e-8)

        return pointacc, pointacc_per_class, voxacc, voxacc_per_class, voxcaliacc, mask

    def compute_miou(self, coords, preds, targets, weights, num_classes):
        """compute miou"""
        coords, preds, targets, weights = self.filter_points(coords, preds, targets, weights)
        seen_classes = np.unique(targets)
        mask = np.zeros(num_classes)
        mask[seen_classes] = 1

        pointmiou = np.zeros(num_classes)
        voxmiou = np.zeros(num_classes)

        uvidx, uvlabel, _ = self.point_cloud_label_to_surface_voxel_label_fast(coords, np.concatenate(
            (np.expand_dims(targets, 1), np.expand_dims(preds, 1)), axis=1), res=0.02)
        for l in seen_classes:
            target_label = np.arange(targets.shape[0])[targets == l]
            pred_label = np.arange(preds.shape[0])[preds == l]
            num_intersection_label = np.intersect1d(pred_label, target_label).shape[0]
            num_union_label = np.union1d(pred_label, target_label).shape[0]
            pointmiou[l] = num_intersection_label / (num_union_label + 1e-8)

            target_label_vox = uvidx[(uvlabel[:, 0] == l)]
            pred_label_vox = uvidx[(uvlabel[:, 1] == l)]
            num_intersection_label_vox = np.intersect1d(pred_label_vox, target_label_vox).shape[0]
            num_union_label_vox = np.union1d(pred_label_vox, target_label_vox).shape[0]
            voxmiou[l] = num_intersection_label_vox / (num_union_label_vox + 1e-8)

        return pointmiou, voxmiou, mask

    def update(self, *inputs):
        """record data"""
        if len(inputs) != 4:
            raise ValueError('Distribute accuracy needs 2 input (y_correct), but got {}'.format(len(inputs)))
        coords = self._convert_data(inputs[0])
        y_pred = self._convert_data(inputs[1])
        target = self._convert_data(inputs[2])
        weights = self._convert_data(inputs[3])

        B, N, chan = coords.shape
        pred = np.expand_dims(y_pred, 0)
        preds = pred.argmax(3)

        coords = coords.reshape(-1, chan)  # (CK*N, C)
        preds = preds.squeeze(0).reshape(-1)  # (CK*N, C)
        targets = target.reshape(-1) # (CK*N, C)
        weights = weights.reshape(-1)  # (CK*N, C)
        pointacc, pointacc_per_class, voxacc, voxacc_per_class, voxcaliacc, acc_mask = self.compute_acc(coords, preds,
                                                                                                   targets, weights,
                                                                                                   self.num_classes)
        pointmiou, voxmiou, miou_mask = self.compute_miou(coords, preds, targets, weights, self.num_classes)
        assert acc_mask.all() == miou_mask.all()
        mask = acc_mask

        # dump
        self.pointacc_list.append(pointacc)
        self.pointacc_per_class_array[self.id] = pointacc_per_class
        self.voxacc_list.append(voxacc)
        self.voxacc_per_class_array[self.id] = voxacc_per_class
        self.voxcaliacc_list.append(voxcaliacc)
        self.pointmiou_per_class_array[self.id] = pointmiou
        self.voxmiou_per_class_array[self.id] = voxmiou
        self.masks[self.id] = mask
        self.id = self.id + 1

    def eval(self):
        # avg_pointacc = np.mean(self.pointacc_list)
        # avg_pointacc_per_class = np.sum(self.pointacc_per_class_array * self.masks, axis=0) / np.sum(self.masks, axis=0)
        # avg_pointacc_per_class = np.nan_to_num(avg_pointacc_per_class)

        # avg_voxacc = np.mean(self.voxacc_list)
        # avg_voxacc_per_class = np.sum(self.voxacc_per_class_array * self.masks, axis=0) / np.sum(self.masks, axis=0)
        # avg_voxacc_per_class = np.nan_to_num(avg_voxacc_per_class)
        # avg_voxcaliacc = np.mean(self.voxcaliacc_list)

        np.seterr(divide='ignore', invalid='ignore')
        avg_pointmiou_per_class = np.sum(self.pointmiou_per_class_array * self.masks, axis=0) / np.sum(self.masks, axis=0)
        avg_pointmiou_per_class = np.nan_to_num(avg_pointmiou_per_class)
        avg_pointmiou = np.mean(avg_pointmiou_per_class)

        # avg_voxmiou_per_class = np.sum(self.voxmiou_per_class_array * self.masks, axis=0) / np.sum(self.masks, axis=0)
        # avg_voxmiou_per_class = np.nan_to_num(avg_voxmiou_per_class)
        # avg_voxmiou = np.mean(avg_voxmiou_per_class)
        return avg_pointmiou

class CustomWithLossCell(nn.Cell):
    """连接前向网络和损失函数"""
    def __init__(self, backbone, loss_fn):
        """输入有两个，前向网络backbone和损失函数loss_fn"""
        super(CustomWithLossCell, self).__init__(auto_prefix=False)
        self._backbone = backbone
        self._loss_fn = loss_fn

    def construct(self, data, label, weights):
        "net construct"
        output = self._backbone(data)  # 前向计算得到网络输出
        return self._loss_fn(ops.Reshape()(output, (-1, 20)), label.view(-1), weights.view(-1))

class ScannetWithEvalCell(nn.Cell):
    """eval-output"""
    def __init__(self, network, add_cast_fp32=False):
        super(ScannetWithEvalCell, self).__init__(auto_prefix=False)
        assert isinstance(add_cast_fp32, bool), "add_cast_fp32 should be float"
        self._network = network
        self.cat = ops.Concat(axis=-1)

    def construct(self, data, label, weights):
        _, N, _ = data.shape
        pred = self._network(data)
        return data, pred, label, weights

class CallbackSaveBymIoU(Callback):
    """SaveCallback"""
    def __init__(self, eval_model, ds_eval, eval_period=1, eval_start=1, save_path=None):
        """init"""
        super(CallbackSaveBymIoU, self).__init__()
        self.model = eval_model
        self.ds_eval = ds_eval
        self.save_path = save_path
        self.eval_period = eval_period
        self.eval_start = eval_start
        self.avg_pointmiou = 0.

    def epoch_end(self, run_context):
        """epoch end"""
        cb_params = run_context.original_args()
        cur_epoch = cb_params.cur_epoch_num
        rank_id = get_rank_id()
        if ((cur_epoch + 1) % self.eval_period) == 0:
            if cur_epoch < self.eval_start:
                return
            if rank_id == 0:
                print("Start evaluate...")
            result = self.model.eval(self.ds_eval)

            avg_pointmiou = result['IoU']
            if avg_pointmiou > self.avg_pointmiou:
                self.avg_pointmiou = avg_pointmiou
                file_name = f"best_model_dev_{rank_id}.ckpt"

                save_path = os.path.join(self.save_path, file_name)
                if not os.path.exists(self.save_path):
                    os.mkdir(self.save_path)
                print("Save model...")
                save_checkpoint(save_obj=cb_params.train_network, ckpt_file_name=save_path)
            print(f"Device:{rank_id}, Epoce:{cur_epoch}, avg_pointmiou:{avg_pointmiou:.5f}")

    def end(self, run_context):
        _ = run_context.original_args()
        rank_id = get_rank_id()
        print(f"Device:{rank_id}, Best Instance mIoU:{(self.ins_mIoU*100):.2f}%, Class mIoU:{(self.cls_mIoU*100):.2f}%")
